/*
 * SSL Socket Factory which performs the MITM cert generation
 *
 * @author Jonathan Walsh 
 * @author Jesse Brown
 */

//Based on SnifferSSLSocketFactory.java from The Grinder distribution.
// The Grinder distribution is available at http://grinder.sourceforge.net/

package mitm;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import javax.net.ServerSocketFactory;
import javax.net.SocketFactory;
import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.SSLServerSocketFactory;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import java.math.BigInteger;

import java.util.Calendar;
import java.io.FileOutputStream;
import iaik.asn1.ObjectID;
import iaik.asn1.structures.*;
import java.util.regex.Pattern;
import java.util.regex.Matcher;
import java.security.PublicKey;
import java.util.Random;
import javax.net.ssl.KeyManager;

/**
 * MITMSSLSocketFactory is used to create SSL sockets.
 *
 * This is needed because the javax.net.ssl socket factory classes don't
 * allow creation of factories with custom parameters.
 *
 */
public final class MITMSSLSocketFactory implements MITMSocketFactory
{
    final ServerSocketFactory m_serverSocketFactory;
    final SocketFactory m_clientSocketFactory;
    final SSLContext m_sslContext;

    public KeyStore ks = null;

    /*
     *
     * We can't install our own TrustManagerFactory without messing
     * with the security properties file. Hence we create our own
     * SSLContext and initialise it. Passing null as the keystore
     * parameter to SSLContext.init() results in a empty keystore
     * being used, as does passing the key manager array obtain from
     * keyManagerFactory.getInstance().getKeyManagers(). To pick up
     * the "default" keystore system properties, we have to read them
     * explicitly. UGLY, but necessary so we understand the expected
     * properties.
     *
     */

    /**
     * This constructor will create an SSL server socket factory
     * that is initialized with a fixed CA certificate
     */
    public MITMSSLSocketFactory()
	throws IOException,GeneralSecurityException
    {
	m_sslContext = SSLContext.getInstance("SSL");

	final KeyManagerFactory keyManagerFactory =
	    KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());

	final String keyStoreFile = System.getProperty(JSSEConstants.KEYSTORE_PROPERTY);
	final char[] keyStorePassword = System.getProperty(JSSEConstants.KEYSTORE_PASSWORD_PROPERTY, "").toCharArray();
	final String keyStoreType = System.getProperty(JSSEConstants.KEYSTORE_TYPE_PROPERTY, "jks");

	final KeyStore keyStore;
	
	if (keyStoreFile != null) {
	    keyStore = KeyStore.getInstance(keyStoreType);
	    keyStore.load(new FileInputStream(keyStoreFile), keyStorePassword);

	    this.ks = keyStore;
	} else {
	    keyStore = null;
	}

	keyManagerFactory.init(keyStore, keyStorePassword);

	m_sslContext.init(keyManagerFactory.getKeyManagers(),
			  new TrustManager[] { new TrustEveryone() },
			  null);

	m_clientSocketFactory = m_sslContext.getSocketFactory();
	m_serverSocketFactory = m_sslContext.getServerSocketFactory(); 
    }

    /**
     * This constructor will create an SSL server socket factory
     * that is initialized with a dynamically generated server certificate
     * that contains the specified common name.
     */
    public MITMSSLSocketFactory(String remoteCN, BigInteger serialno)
	throws IOException,GeneralSecurityException, Exception {
	
	// Load in the default values for this session
	m_sslContext = SSLContext.getInstance("SSL");

	final String keyStoreFile = System.getProperty(JSSEConstants.KEYSTORE_PROPERTY);
	final char[] keyStorePassword = System.getProperty(JSSEConstants.KEYSTORE_PASSWORD_PROPERTY, "").toCharArray();
	final String keyStoreType = System.getProperty(JSSEConstants.KEYSTORE_TYPE_PROPERTY, "jks");
        final KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());

	// if the keystore file isn't passed in we have a problem
	if ( keyStoreFile == null ) {
	    System.err.println( "Failed to read in keystore file!" );
	    System.err.println( "Must pass this value in!" );
	    System.exit(1);
	}

	// make the key store for this session
        final KeyStore keyStore;
	keyStore = KeyStore.getInstance(keyStoreType);
	keyStore.load(new FileInputStream(keyStoreFile), keyStorePassword);
	this.ks = keyStore;
		
	// Get the private key we want to use
	PrivateKey privKey = (PrivateKey)ks.getKey( "mykey", 
						    keyStorePassword );
	
	// Get the public key
	X509Certificate cert = (X509Certificate)ks.getCertificate("mykey");
	PublicKey pubKey = (PublicKey)cert.getPublicKey();
	
	// Construct the issuerDN from the actual server's cert.  We need 
	// to do this so that the client can't tell that we aren't the end 
	// point they are looking for
	String issuerDNString = cert.getIssuerDN().getName();
	String issuerCN = "";
	Pattern commonNamePattern = Pattern.compile("CN=([^,]+)");
	Matcher m = commonNamePattern.matcher(issuerDNString);
	if (m.find()) issuerCN = m.group(1);
	Name issuerDN = new Name();
	issuerDN.addRDN(ObjectID.commonName, issuerCN);
	Name subjectDN = new Name();
	subjectDN.addRDN(ObjectID.commonName, remoteCN);
	
	// Create the new fake certificate - Ok seriously why does this 
	// have to be the same name as the Java one...
	iaik.x509.X509Certificate evilCert = new iaik.x509.X509Certificate();
	
	// Set my fake DN (which looks like the endpoint)
	evilCert.setIssuerDN(issuerDN);
	evilCert.setSubjectDN(subjectDN);
	
	// Serial number can be anything, this is all fake anyway
	evilCert.setSerialNumber(new BigInteger(32, new Random()));
	
	// Fill in with some valid-ish dates
	Calendar cal = Calendar.getInstance();
	evilCert.setValidNotBefore( cal.getTime() );
	cal.add(Calendar.MONTH, 12); // 12 months is enough for anybody!
	evilCert.setValidNotAfter( cal.getTime() );
	
	// Set the keys and sign this cert
	evilCert.setPublicKey(pubKey);
	evilCert.sign( AlgorithmID.sha1WithRSAEncryption, privKey );
	
	// Add the cert to the keystore
	ks.setKeyEntry( remoteCN, 
			privKey, 
			keyStorePassword, 
			new iaik.x509.X509Certificate[] { evilCert } );
	FileOutputStream out = new FileOutputStream( keyStoreFile );
	ks.store( out,//new FileOutputStream( keyStoreFile ),
		  keyStorePassword );
	out.close();

	// init the actual factory with either the keystore we just made, or 
	//   the one we are reusing
        keyManagerFactory.init(ks, keyStorePassword);
        KeyManager[] keys = keyManagerFactory.getKeyManagers();

        m_sslContext.init(keys, new TrustManager[] { new TrustEveryone() }, null);

        m_clientSocketFactory = m_sslContext.getSocketFactory();
        m_serverSocketFactory = m_sslContext.getServerSocketFactory();
    }


    public final ServerSocket createServerSocket(String localHost,
						 int localPort,
						 int timeout)
	throws IOException
    {
	final SSLServerSocket socket =
	    (SSLServerSocket)m_serverSocketFactory.createServerSocket(
		localPort, 50, InetAddress.getByName(localHost));

	socket.setSoTimeout(timeout);

	socket.setEnabledCipherSuites(socket.getSupportedCipherSuites());

	return socket;
    }

    public final Socket createClientSocket(String remoteHost, int remotePort)
	throws IOException
    {
	final SSLSocket socket =
	    (SSLSocket)m_clientSocketFactory.createSocket(remoteHost,
							  remotePort);

	socket.setEnabledCipherSuites(socket.getSupportedCipherSuites());
	
	socket.startHandshake();

	return socket;
    }

    /**
     * We're carrying out a MITM attack, we don't care whether the cert
     * chains are trusted or not ;-)
     *
     */
    private static class TrustEveryone implements X509TrustManager
    {
	public void checkClientTrusted(X509Certificate[] chain,
				       String authenticationType) {
	}
	
	public void checkServerTrusted(X509Certificate[] chain,
				       String authenticationType) {
	}

	public X509Certificate[] getAcceptedIssuers()
	{
	    return null;
	}
    }
}
    
